// 此代码由DeepSeek生成的代码修改而来
#include "eigen/iterator.h"

#include <random>
#include <vector>
#include <algorithm>

namespace  ldl_eigen
{
Iterator::Iterator(int64_t data_size) : data_size_(data_size)
{
    // 初始化索引向量
    indices_.resize(data_size_);
    for (int64_t i = 0; i < data_size_; ++i)
    {
        indices_[i] = i;
    }
    
    // 使用随机设备生成种子
    std::random_device rd;
    rng_.seed(rd());

    reset();
}


Eigen::MatrixXf Iterator::get_rand_batch(const Eigen::MatrixXf &data, int64_t batch)
{
    // 检查批次大小是否有效
    if (batch <= 0 || batch > data_size_)
    {
        throw std::invalid_argument("Invalid batch size");
    }
    
    // 如果当前索引加上批次大小超过数据大小，重新打乱索引
    if (current_idx_ + batch > data_size_)
    {
        reset();
    }
    
    // 创建批次矩阵
    Eigen::MatrixXf batch_data(batch, data.cols());
    
    // 填充批次数据
    for (int64_t i = 0; i < batch; ++i)
    {
        batch_data.row(i) = data.row(indices_[current_idx_]);
        ++current_idx_;
    }
    
    return batch_data;
}

void Iterator::get_rand_batch(const Eigen::MatrixXf &feature, const Eigen::MatrixXf &label, 
        Eigen::MatrixXf &feature_batch, Eigen::MatrixXf &label_batch, int64_t batch_size)
{
    // 1. 参数检查
    if ((feature.size() == 0) || (label.size() == 0) || feature.rows() != label.rows()) {
        throw std::invalid_argument("Feature and label sizes must match");
    }
    if (batch_size <= 0) {
        throw std::invalid_argument("Batch size must be positive");
    }

    // 3. 检查是否需要重新洗牌
    if (current_idx_ + batch_size > data_size_)
    {
        reset();
    }

    // 4. 预分配批次内存
    // feature_batch.resize(batch_size, feature.cols());
    // label_batch.resize(batch_size, label.cols());
    feature_batch = Eigen::MatrixXf(batch_size, feature.cols());
    label_batch =  Eigen::MatrixXf(batch_size, label.cols());

    // 5. 填充批次数据
    std::vector<int> ind(indices_.begin() + current_idx_, indices_.begin() + current_idx_+batch_size);
    current_idx_ += batch_size;
    feature_batch = feature(ind, Eigen::all);
    label_batch = label(ind, Eigen::all);
}

void Iterator::reset()
{
    // 打乱索引顺序
    std::shuffle(indices_.begin(), indices_.end(), rng_);
    
    // 重置当前索引位置
    current_idx_ = 0;
}

}